package myssh

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"more_ssh/conf"
	"more_ssh/middleware"
	"more_ssh/sql"
	"more_ssh/util"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"golang.org/x/crypto/ssh"
)

var (
	// sz fmt.Sprintf("%+q", "rz\r**\x18B00000000000000\r\x8a\x11")
	//ZModemSZStart = []byte{13, 42, 42, 24, 66, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 13, 138, 17}
	ZModemSZStart = []byte{42, 42, 24, 66, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 13, 138, 17}
	// sz 结束 fmt.Sprintf("%+q", "\r**\x18B0800000000022d\r\x8a")
	//ZModemSZEnd = []byte{13, 42, 42, 24, 66, 48, 56, 48, 48, 48, 48, 48, 48, 48, 48, 48, 50, 50, 100, 13, 138}
	ZModemSZEnd = []byte{42, 42, 24, 66, 48, 56, 48, 48, 48, 48, 48, 48, 48, 48, 48, 50, 50, 100, 13, 138}
	// sz 结束后可能还会发送两个 OO，但是经过测试发现不一定每次都会发送 fmt.Sprintf("%+q", "OO")
	ZModemSZEndOO = []byte{79, 79}

	// rz fmt.Sprintf("%+q", "**\x18B0100000023be50\r\x8a\x11")
	ZModemRZStart = []byte{42, 42, 24, 66, 48, 49, 48, 48, 48, 48, 48, 48, 50, 51, 98, 101, 53, 48, 13, 138, 17}
	// rz -e fmt.Sprintf("%+q", "**\x18B0100000063f694\r\x8a\x11")
	ZModemRZEStart = []byte{42, 42, 24, 66, 48, 49, 48, 48, 48, 48, 48, 48, 54, 51, 102, 54, 57, 52, 13, 138, 17}
	// rz -S fmt.Sprintf("%+q", "**\x18B0100000223d832\r\x8a\x11")
	ZModemRZSStart = []byte{42, 42, 24, 66, 48, 49, 48, 48, 48, 48, 48, 50, 50, 51, 100, 56, 51, 50, 13, 138, 17}
	// rz -e -S fmt.Sprintf("%+q", "**\x18B010000026390f6\r\x8a\x11")
	ZModemRZESStart = []byte{42, 42, 24, 66, 48, 49, 48, 48, 48, 48, 48, 50, 54, 51, 57, 48, 102, 54, 13, 138, 17}
	// rz 结束 fmt.Sprintf("%+q", "**\x18B0800000000022d\r\x8a")
	ZModemRZEnd = []byte{42, 42, 24, 66, 48, 56, 48, 48, 48, 48, 48, 48, 48, 48, 48, 50, 50, 100, 13, 138}

	// **\x18B0
	ZModemRZCtrlStart = []byte{42, 42, 24, 66, 48}
	// \r\x8a\x11
	ZModemRZCtrlEnd1 = []byte{13, 138, 17}
	// \r\x8a
	ZModemRZCtrlEnd2 = []byte{13, 138}

	// zmodem 取消 \x18\x18\x18\x18\x18\x08\x08\x08\x08\x08
	ZModemCancel = []byte{24, 24, 24, 24, 24, 8, 8, 8, 8, 8}
)

// 定义一个结构体 方便保存各种连接信息
type MySSH struct {
	Websocket *websocket.Conn
	Stdin     io.WriteCloser
	// Stdout    *wsBufferWriter
	Stdout    io.Reader
	Stderr    io.Reader
	Session   *ssh.Session
	SshClient *ssh.Client

	ZModemSZ, ZModemRZ, ZModemSZOO bool
}

type WebMsg struct {
	Type string `json:"type"`
	Cols int    `json:"Cols"`
	Rows int    `json:"rows"`
	Data []byte `json:"data"`
	// ReslutData string `json:"reslutData"`
}

// 程序入口
func RunWebSSH(c *gin.Context) {
	authorization, ok := c.GetQuery("Authorization")
	if !ok {
		fmt.Println("authorization获取失败")
		return
	}
	// 校验token
	userId, ok := middleware.WsParseToken(authorization)
	if !ok {
		fmt.Println("token解析失败")
		return
	}

	audit := &sql.Audit{}

	//获取前端传入的参数
	id, ok := c.GetQuery("id")
	if !ok {
		fmt.Println("失败")
		return
	}
	idInt, err := strconv.Atoi(id)
	if err != nil {
		fmt.Println("传的参数不是整数")
		return
	}

	account, ok := c.GetQuery("account")
	if !ok {
		fmt.Println("获取account失败")
	} else {
		fmt.Println(account)
	}

	host := &sql.Host{
		ID:      uint(idInt),
		Account: account,
	}

	conf.DB.First(host)
	if host == nil {
		fmt.Println("没查到id=" + id)
		return
	}

	ip := host.IP
	port := host.Port

	addr := fmt.Sprintf("%s:%d", ip, port)

	fmt.Println(addr)

	accountInfo := &sql.Account{}
	conf.DB.Where("account=? AND host_id=?", account, id).First(accountInfo)
	fmt.Println(accountInfo.Password)

	//记录审计相关信息
	audit.StartTime = int(time.Now().Unix())
	audit.Account = account
	// audit.UserId = c.MustGet("userId").(uint)
	audit.HostId = uint(idInt)
	fileName := "log/" + util.RandString(10) + "-20220107-admin"
	audit.FilePath = fileName
	audit.ClientIp = c.ClientIP()
	audit.Service = "ssh"
	audit.UserId = userId

	conf.DB.Create(audit)

	// 创建读写文件
	fileTimeLine, err := os.OpenFile(fileName+".tl", os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		fmt.Println("创建fileTimeLine文件失败")
		return
	}
	fileOperator, err := os.OpenFile(fileName+".data", os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		fmt.Println("创建fileOperator文件失败")
		return
	}
	mySSH := &MySSH{}

	// 1. 升级请求websocket
	upGrader := websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024 * 1024 * 10,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	webcon, err := upGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		fmt.Println("升级http 为websoket失败：", err)
	}

	defer webcon.Close()

	mySSH.Websocket = webcon // 将websocket连接保存到对象中

	// 创建一个ssh的配置
	config := &ssh.ClientConfig{
		Timeout:         time.Second * 15, //ssh 连接time out 时间一秒钟, 如果ssh验证错误 会在一秒内返回
		User:            accountInfo.Account,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), //这个可以， 但是不够安全
		//HostKeyCallback: hostKeyCallBackFunc(h.Host),
		Auth: []ssh.AuthMethod{ssh.Password(accountInfo.Password)},
		Config: ssh.Config{
			Ciphers: []string{"aes128-ctr", "aes192-ctr", "aes256-ctr", "aes128-gcm@openssh.com", "arcfour256", "arcfour128", "aes128-cbc", "3des-cbc", "aes192-cbc", "aes256-cbc"},
		},
	}

	// 创建一个客户端
	sshClient, err := ssh.Dial("tcp", addr, config)
	// sshClient, err := ssh.Dial("tcp", "192.168.2.72:22", config)
	if err != nil {
		fmt.Println(err)
		return
	}

	defer sshClient.Close()
	mySSH.SshClient = sshClient

	session, err := sshClient.NewSession()
	if err != nil {
		fmt.Println(err)
		return
	}
	defer session.Close()
	mySSH.Session = session

	// 保存输入流
	mySSH.Stdin, err = session.StdinPipe()
	if err != nil {
		fmt.Println(err)
		return
	}

	//保存ssh输出流
	// sshOut := new(wsBufferWriter)
	// session.Stdout = sshOut
	// session.Stderr = sshOut
	mySSH.Stdout, err = session.StdoutPipe()
	if err != nil {
		fmt.Println("获取输出通道失败：", err)
		return
	}

	modes := ssh.TerminalModes{
		ssh.ECHO:          1,     // disable echo
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}
	// Request pseudo terminal
	if err := session.RequestPty("xterm", 30, 120, modes); err != nil {
		fmt.Println("绑定pty失败：", err)
		return
	}

	session.Shell()

	quitChan := make(chan bool, 3)

	//执行远程命令
	go Send2SSH(mySSH, quitChan)
	go Send2Web(mySSH, fileTimeLine, fileOperator, quitChan)
	go SshWait(mySSH, quitChan)
	<-quitChan
	audit.EndTime = int(time.Now().Unix())
	conf.DB.Save(audit)
	fmt.Println("关闭会话")

}

// 读取websocket数据，发送到ssh输入流中
func Send2SSH(mySSh *MySSH, quitChan chan bool) {
	defer setQuitChan(quitChan)
	defer fmt.Println("退出Send2SSH")
	for {
		//read websocket msg  需要通过msgType 判断是传输类型
		dataType, wsData, err := mySSh.Websocket.ReadMessage()
		if err != nil {
			fmt.Println("读取websocket数据失败：", err)
			return
		}
		//二进制数据 直接发送到ssh通道
		if dataType == websocket.BinaryMessage {
			mySSh.Stdin.Write(wsData)
			continue
		}

		msg := &WebMsg{}
		err = json.Unmarshal(wsData, msg)

		if err != nil {
			fmt.Println("解析数据转化成json失败：", err)
			msg.Type = "cmd" // return
		}
		switch msg.Type {
		case "resize":
			err = mySSh.Session.WindowChange(msg.Rows, msg.Cols)
			if err != nil {
				fmt.Println("重设宽高失败：", err)
			}
		case "stdin":
			data1, _ := url.QueryUnescape(string(msg.Data))
			_, err = mySSh.Stdin.Write([]byte(data1))
			// fmt.Println("收到数据", string(msg.Data))
			if err != nil {
				fmt.Println("ssh发送数据失败：", err)
			}

		case "cmd":
			_, err = mySSh.Stdin.Write(wsData)
			if err != nil {
				fmt.Println("cmd  ssh发送数据失败：", err)
			}

		case "ignore":
			// fmt.Println(string(msg.Data))

		}

	}

}

// 读取ssh输出，发送到websocket中
func Send2Web(mySSh *MySSH, fileTimeLine *os.File, fileOperator *os.File, quitChan chan bool) {
	defer setQuitChan(quitChan)
	defer fmt.Println("退出Send2Web")
	for {
		// n := mySSh.Stdout.buffer.Len()
		sshOut := make([]byte, 8192)
		n, err := mySSh.Stdout.Read(sshOut)
		if err != nil {
			return
		}

		if n > 0 {
			//拿到 sshOut输出后，首先判断当前文件传输状态是否结束
			if mySSh.ZModemSZOO {
				// 经过测试 centos7-8 使用的 lrzsz-0.12.20 在 sz 结束时会发送 ZModemSZEndOO
				// 而 deepin20 等自带更新的 lrzsz-0.12.21rc 在 sz 结束时不会发送 ZModemSZEndOO， 而前端 zmodemjs
				// 库只有接收到 ZModemSZEndOO 才会认为 sz 结束，固这里需判断 sz 结束时是否发送了 ZModemSZEndOO，
				// 如果没有则手动发送一个，以便保证前端 zmodemjs 库正常运行（如果不发送，会导致使用 sz 命令时无法连续
				// 下载多个文件）。
				mySSh.ZModemSZOO = false
				if n < 2 {
					// 手动发送 ZModemSZEndOO
					mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemSZEndOO)
					mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", Data: sshOut[:n]})
					// mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", ReslutData: string(sshOut[:n])})
				} else if n == 2 {
					if sshOut[0] == ZModemSZEndOO[0] && sshOut[1] == ZModemSZEndOO[1] {
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemSZEndOO)
					} else {
						// 手动发送 ZModemSZEndOO
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemSZEndOO)
						mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", Data: sshOut[:n]})
						// mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", ReslutData: string(sshOut[:n])})

					}
				} else {
					if sshOut[0] == ZModemSZEndOO[0] && sshOut[1] == ZModemSZEndOO[1] {
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, sshOut[:2])
						mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", Data: sshOut[2:n]})
						// mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", ReslutData: string(sshOut[:n])})

					} else {
						// 手动发送 ZModemSZEndOO
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemSZEndOO)
						mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", Data: sshOut[:n]})
						// mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", ReslutData: string(sshOut[:n])})

					}
				}

			} else {
				//判断当前状态是否是文件上传或下载， 再判断包有没有文件传输的信息
				if mySSh.ZModemSZ {
					//判断是否有结束标记
					if x, ok := IsContain(sshOut[:n], ZModemSZEnd); ok { //判断是完成下载
						mySSh.ZModemSZ = false
						mySSh.ZModemSZOO = true
						fmt.Println("结束下载")
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemSZEnd)
						if len(x) != 0 {
							mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: x}) //原来是console打印
							// mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", ReslutData: string(sshOut[:n])})
						}
					} else if _, ok := IsContain(sshOut[:n], ZModemCancel); ok { //判断是否取消下载
						fmt.Println("取消下载")
						mySSh.ZModemSZ = false
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, sshOut[:n])
					} else { //都没有直接发送下载的数据
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, sshOut[:n])
					}

				} else if mySSh.ZModemRZ {
					// fmt.Println("rz")
					if x, ok := IsContain(sshOut[:n], ZModemRZEnd); ok {
						fmt.Println(string(x))
						mySSh.ZModemRZ = false
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemRZEnd)
						if len(x) != 0 {
							mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: x})

						}
					} else if _, ok := IsContain(sshOut[:n], ZModemCancel); ok {
						mySSh.ZModemRZ = false
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, sshOut[:n])
					} else {
						// rz 上传过程中服务器端还是会给客户端发送一些信息，比如心跳
						//ws.Ws.WriteJSON(&message{Type: messageTypeConsole, Data: sshOut[:n]})
						//ws.Ws.WriteMessage(websocket.BinaryMessage, sshOut[:n])
						startIndex := bytes.Index(sshOut[:n], ZModemRZCtrlStart)
						if startIndex != -1 {
							endIndex := bytes.Index(sshOut[:n], ZModemRZCtrlEnd1)
							if endIndex != -1 {
								ctrl := append(ZModemRZCtrlStart, sshOut[startIndex+len(ZModemRZCtrlStart):endIndex]...)
								ctrl = append(ctrl, ZModemRZCtrlEnd1...)
								mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ctrl)
								info := append(sshOut[:startIndex], sshOut[endIndex+len(ZModemRZCtrlEnd1):n]...)
								if len(info) != 0 {
									mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: info})
								}
							} else {
								endIndex = bytes.Index(sshOut[:n], ZModemRZCtrlEnd2)
								if endIndex != -1 {
									ctrl := append(ZModemRZCtrlStart, sshOut[startIndex+len(ZModemRZCtrlStart):endIndex]...)
									ctrl = append(ctrl, ZModemRZCtrlEnd2...)
									mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ctrl)
									info := append(sshOut[:startIndex], sshOut[endIndex+len(ZModemRZCtrlEnd2):n]...)
									if len(info) != 0 {
										mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: info})
									}
								} else {
									mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: sshOut[:n]})
								}
							}
						} else {
							mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: sshOut[:n]})
						}
					}
				} else {
					if x, ok := IsContain(sshOut[:n], ZModemSZStart); ok { //是否包含sz

						if y, ok := IsContain(x, ZModemCancel); ok {
							// 下载不存在的文件以及文件夹(zmodem 不支持下载文件夹)时
							mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", Data: y})
							// mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", ReslutData: string(y)})

						} else {
							mySSh.ZModemSZ = true
							if len(x) != 0 {
								mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: x})
								// mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", ReslutData: string(x)})
							}
							mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemSZStart)
						}

					} else if x, ok := IsContain(sshOut[:n], ZModemRZStart); ok {

						mySSh.ZModemRZ = true
						if len(x) != 0 {
							mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: x})
						}
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemRZStart)

					} else if x, ok := IsContain(sshOut[:n], ZModemRZEStart); ok {

						mySSh.ZModemRZ = true
						if len(x) != 0 {
							mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: x})
							// mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", ReslutData: string(x)})
						}
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemRZEStart)

					} else if x, ok := IsContain(sshOut[:n], ZModemRZSStart); ok {

						mySSh.ZModemRZ = true
						if len(x) != 0 {
							mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: x})
							// mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", ReslutData: string(x)})
						}
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemRZSStart)

					} else if x, ok := IsContain(sshOut[:n], ZModemRZESStart); ok {

						mySSh.ZModemRZ = true
						if len(x) != 0 {
							mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", Data: x})
							// mySSh.Websocket.WriteJSON(&WebMsg{Type: "console", ReslutData: string(x)})
						}
						mySSh.Websocket.WriteMessage(websocket.BinaryMessage, ZModemRZESStart)

					} else {
						// 如果不满足上面所有的条件 说明只是普通的命令回执显示，直接发回前	端即可
						mySSh.Websocket.WriteJSON(&WebMsg{Type: "stdin", Data: sshOut[:n]})
						// 将操作记录到文本中
						// 写入操作的时间线
						fileTimeLine.Write([]byte(fmt.Sprintf("%d\t%d\n", time.Now().Unix(), n)))
						//写入操作的内容
						fileOperator.Write(sshOut[:n])
					}
				}
			}
		}
	}
}

//判断切片x 是否包含y ，包含则返回去除y后的x，否则返回false
func IsContain(x, y []byte) (n []byte, contain bool) {
	index := bytes.Index(x, y)
	if index == -1 {
		return n, false
	}
	lastIndex := index + len(y)
	n = append(x[:index], x[lastIndex:]...)
	return n, true
}

//监听ssh会话是否关闭，如果关闭，则关闭所有的连接
func SshWait(mySSH *MySSH, quitChan chan bool) {
	defer setQuitChan(quitChan)
	if err := mySSH.Session.Wait(); err != nil {
		fmt.Println("等待会话退出失败：", err)
		return
	}

}

func setQuitChan(quitChan chan bool) {
	quitChan <- true
}
